package sftp

import (
	"io"
	"os"
	"path"

	"github.com/conduitio/bwlimit"
)

// UploadFile 上传文件，默认sftp配置，不限速
func (c *Client) UploadFile(localFile, remoteDir string) error {
	return c.UploadFileLimit(localFile, remoteDir, 0)
}

// UploadFileWithConf 上传文件，自定义sftp配置，不限速
func (c *Client) UploadFileWithConf(conf *Conf, localFile, remoteDir string) error {
	return c.UploadFileLimitWithConf(conf, localFile, remoteDir, 0)
}

// UploadFileLimit 上传文件，默认sftp配置，限速
func (c *Client) UploadFileLimit(localFile, remoteDir string, rate int) error {
	reader, err := os.Open(localFile)
	if err != nil {
		return err
	}
	defer reader.Close()

	return c.UploadStreamLimit(localFile, remoteDir, reader, rate)
}

// UploadFileLimitWithConf 上传文件，自定义sftp配置，限速
func (c *Client) UploadFileLimitWithConf(conf *Conf, localFile, remoteDir string, rate int) error {
	reader, err := os.Open(localFile)
	if err != nil {
		return err
	}
	defer reader.Close()

	return c.UploadStreamLimitWithConf(conf, localFile, remoteDir, reader, rate)
}

// UploadStreamLimit 上传文件，默认sftp配置，限速
func (c *Client) UploadStreamLimit(localFile, remoteDir string, fileReader io.Reader, rate int) error {
	return c.UploadStreamLimitWithConf(c.defaultConf, localFile, remoteDir, fileReader, rate)
}

// UploadStreamLimitWithConf 上传文件，自定义sftp配置，限速
func (c *Client) UploadStreamLimitWithConf(conf *Conf, localFile, remoteDir string, fileReader io.Reader, rate int) error {
	var (
		err      error
		sftpConn *Conn
	)

	// 连接sftp
	sftpConn, err = c.ConnectWithConf(conf)
	if err != nil {
		return err
	}
	defer sftpConn.Close()

	err = sftpConn.MkdirAll(remoteDir)
	if err != nil {
		return err
	}
	fileName := path.Base(localFile)
	writer, err := sftpConn.Create(path.Join(remoteDir, fileName))
	if err != nil {
		return err
	}
	defer writer.Close()
	// 限速 KB
	writeLimit := bwlimit.Byte(rate) * bwlimit.KB
	// 上传写限速
	uploadWriter := bwlimit.NewWriter(writer, writeLimit)
	defer uploadWriter.Close()

	// 限速写数据
	_, err = io.Copy(uploadWriter, fileReader)
	if err != nil {
		return err
	}
	return nil
}
